(* 1. 定义类型 — 只定义一次！ *)
type ('a, 'b) labeled_graph = {
  nodes : 'a list;
  labeled_edges : ('a * 'a * 'b) list
}

(* 2. 定义图 g *)
let g = {
  nodes = ['a'; 'b'; 'c'; 'd'; 'e'; 'f'; 'g'; 'h'];
  labeled_edges = [
    ('a', 'b', 5); ('a', 'd', 3); ('b', 'c', 2);
    ('b', 'e', 4); ('c', 'e', 6); ('d', 'e', 7);
    ('d', 'f', 4); ('d', 'g', 3); ('e', 'h', 5);
    ('f', 'g', 4); ('g', 'h', 1)
  ]
}

(* 3. 辅助函数：mem *)
let rec mem x = function
  | [] -> false
  | y :: ys -> x = y || mem x ys

(* 4. Prim's algorithm 实现 *)
let ms_tree (g : ('a, 'b) labeled_graph) (cmp : 'b -> 'b -> int) : ('a, 'b) labeled_graph =
  match g.nodes with
  | [] -> { nodes = []; labeled_edges = [] }
  | start_node :: _ ->
      let rec prim visited mst_edges candidates =
        if List.length visited = List.length g.nodes then
          { nodes = g.nodes; labeled_edges = List.sort_uniq (fun (a,b,_) (c,d,_) -> compare (min a b, max a b) (min c d, max c d)) mst_edges }
        else
          let sorted_candidates = List.sort (fun (_, _, w1) (_, _, w2) -> cmp w1 w2) candidates in
          match sorted_candidates with
          | [] -> { nodes = g.nodes; labeled_edges = mst_edges }
          | (u, v, w) :: rest ->
              if mem u visited && not (mem v visited) then
                let new_visited = v :: visited in
                let new_mst = (u, v, w) :: mst_edges in
                let new_edges = List.filter (fun (x, y, _) ->
                    (x = v && not (mem y new_visited)) ||
                    (y = v && not (mem x new_visited))
                ) g.labeled_edges in
                prim new_visited new_mst (new_edges @ rest)
              else if mem v visited && not (mem u visited) then
                let new_visited = u :: visited in
                let new_mst = (u, v, w) :: mst_edges in
                let new_edges = List.filter (fun (x, y, _) ->
                    (x = u && not (mem y new_visited)) ||
                    (y = u && not (mem x new_visited))
                ) g.labeled_edges in
                prim new_visited new_mst (new_edges @ rest)
              else
                prim visited mst_edges rest
      in
      let initial_visited = [start_node] in
      let initial_candidates =
        List.filter (fun (u, v, _) ->
            (u = start_node && v <> start_node) ||
            (v = start_node && u <> start_node)
        ) g.labeled_edges
      in
      prim initial_visited [] initial_candidates;;

(* 5. 比较函数 *)
let int_cmp = compare;;

(* 6. 调用 *)
let mst = ms_tree g int_cmp;;